Pure-python masked UDFs #9174
Pure-python masked UDFs #9174
Conversation
brandon-b-miller
added
feature request
brandon-b-miller
requested review from
trxcllnt,
jrhemstad and
galipremsagar
github-actions
bot
added
CMake
|
cc @gmarkall |
|
Upon consulting with @shwina we decided it was better to leave the c++ alone in this PR and open a separate one removing the c++ side of things. In that case it could be reverted separately without having to revert the new pipeline as well. |
|
This should be ready to go. |
| @@ -11,7 +11,7 @@ | |||
|
|
|||
| from cudf.core.udf.typing import MaskedType, NAType | |||
|
|
|||
| from . import classes | |||
| from . import api | |||
There was a problem hiding this comment.
Can we make these imports absolute to be consistent with the rest of the code-base?
python/cudf/cudf/core/udf/typing.py
Outdated
| from . import api | ||
| from ._ops import arith_ops, comparison_ops |
There was a problem hiding this comment.
Absolute imports here too.
There was a problem hiding this comment.
I think this is now looking pretty good - there are some small suggestions on the diff but they're fairly minor. I wanted to double-check the suggestions made sense so I had to implement them - if you like them, you could just pull in the last three commits from https://github.com/gmarkall/cudf/tree/masked-udfs-20210927 which implement the suggestions on the diff here.
| *(col.mask is None for col in df._data.values()), | ||
| ) | ||
| if precompiled.get(cache_key) is not None: | ||
| kernel, scalar_return_type = precompiled[cache_key] |
There was a problem hiding this comment.
Could you just return precompiled[cache_key] here to save the rest of the function living in an else block?
| # The python function definition representing the kernel | ||
| _kernel = local_exec_context["_kernel"] | ||
| kernel = cuda.jit(sig)(_kernel) | ||
| precompiled[cache_key] = (kernel, scalar_return_type) |
There was a problem hiding this comment.
Could you cache (kernel, numpy_support.as_dtype(scalar_return_type) so that you don't need to call as_dtype on the scalar_return_type each time it's returned?
|
rerun tests |
| cvarbytes = b"" | ||
|
|
||
| key = (type_signature, codebytes, cvarbytes) | ||
| key = make_cache_key(udf, type_signature) |
There was a problem hiding this comment.
Let's also move the 2 comments above into the make_cache_key function instead, as a "docstring".
|
@gpucibot merge |
|
This is awesome work, @brandon-b-miller! |
Depends on #9174 Adds `Series.apply` which applies a scalar UDF elementwise to the series data returning a new series. Null sensitive. Works in terms of our numba `MaskedType` extension type. Similar to `pd.Series.apply`. Authors: - https://github.com/brandon-b-miller Approvers: - Ashwin Srinath (https://github.com/shwina) - H. Thomson Comer (https://github.com/thomcom) - GALI PREM SAGAR (https://github.com/galipremsagar) URL: #9217
This PR removes the c++ side of the original masked UDF code introduced in #8213. These kernels had some limitations and are now superseded by the numba-generated versions we moved to in #9174. As far as I can tell, cuDF python was the only thing consuming this API for the short time it has existed. However I am marking this breaking just in case. Authors: - https://github.com/brandon-b-miller Approvers: - Mark Harris (https://github.com/harrism) - David Wendt (https://github.com/davidwendt) - Vyas Ramasubramani (https://github.com/vyasr) URL: #9792
Replaces C++ implementation of masked UDF pipeline with a pure python version which compiles and launches the entire kernel using numba. This solves a bunch of problems:
cuda::std::tupleto work with NVRTC 11.0.pow,sin, andcosis now provided since all the PTX is compiled and linked inside numba (Fixes [FEA] Support UDF Runtime compilation for incoming PTX with non-inlineable callees #8470)